#pragma once
#include "trt_define.h"
#include <string>
#include <memory>
#include <cuda_runtime.h>
#include "logger.h"

_TRT_INFER_BEGIN
#define GPU_BLOCK_THREADS 512
template <typename _T>
std::shared_ptr<_T> make_nvshared(_T *ptr)
{
    return std::shared_ptr<_T>(ptr, [](_T *p)
                               { p->destroy(); });
}

#define checkRuntime(call) CUDATools::checkCudaRuntime(call, #call, __LINE__, __FILE__)

#define checkKernel(...)                                                 \
    __VA_ARGS__;                                                         \
    do                                                                   \
    {                                                                    \
        cudaError_t cudaStatus = cudaPeekAtLastError();                  \
        if (cudaStatus != cudaSuccess)                                   \
        {                                                                \
            LOGE << "launch failed: " << cudaGetErrorString(cudaStatus); \
        }                                                                \
    } while (0);

#define checkCudaKernel(...)                                             \
    __VA_ARGS__;                                                         \
    do                                                                   \
    {                                                                    \
        cudaError_t cudaStatus = cudaPeekAtLastError();                  \
        if (cudaStatus != cudaSuccess)                                   \
        {                                                                \
            LOGE << "launch failed: " << cudaGetErrorString(cudaStatus); \
        }                                                                \
    } while (0);

namespace CUDATools
{
    bool checkCudaRuntime(cudaError_t e, const char *call, int iLine, const char *szFile);
    bool check_device_id(int device_id);
    int current_device_id();
    std::string device_description();
    void set_device_id(int device_id);

    dim3 grid_dims(int numJobs);
    dim3 block_dims(int numJobs);
    inline int upbound(int n, int align = 32) { return (n + align - 1) / align * align; }

    class AutoDevice
    {
    public:
        AutoDevice(int device_id = 0);
        virtual ~AutoDevice();

    private:
        int old_ = -1;
    };
}
_TRT_INFER_END